Add lax.tile_p

PiperOrigin-RevId: 825382043

Author: sharadmv

jax/_src/lax/lax.py

@@ -6729,6 +6729,101 @@ def _broadcast_in_dim_ragged_prop_rule(eqn_params, invar_raggedness, outvars):
 )
 
 
+def tile(operand: ArrayLike, reps: Sequence[int]) -> Array:
+  """Tiles an array by repeating it along each dimension.
+
+  Args:
+    operand: an array to tile.
+    reps: a sequence of integers representing the number of repeats for each
+      dimension. Must have the same length as ``operand.ndim``.
+
+  Returns:
+    A tiled array with shape ``(operand.shape[0] * reps[0], ...,
+    operand.shape[-1] * reps[-1])``.
+
+  Examples:
+    >>> x = jnp.array([[1, 2], [3, 4]])
+    >>> lax.tile(x, (2, 3))
+    Array([[1, 2, 1, 2, 1, 2],
+           [3, 4, 3, 4, 3, 4],
+           [1, 2, 1, 2, 1, 2],
+           [3, 4, 3, 4, 3, 4]], dtype=int32)
+
+    >>> y = jnp.array([1, 2, 3])
+    >>> lax.tile(y, (2,))
+    Array([1, 2, 3, 1, 2, 3], dtype=int32)
+
+    >>> z = jnp.array([[1], [2]])
+    >>> lax.tile(z, (1, 3))
+    Array([[1, 1, 1],
+           [2, 2, 2]], dtype=int32)
+  """
+  return tile_p.bind(operand, reps=tuple(reps))
+
+
+def _tile_abstract_eval(operand, *, reps):
+  if len(reps) != operand.ndim:
+    raise ValueError(
+        'tile reps must have length equal to operand.ndim, '
+        f'got reps={reps} for operand.ndim={operand.ndim}'
+    )
+  out_shape = tuple(d * r for d, r in zip(operand.shape, reps))
+  return operand.update(shape=out_shape)
+
+
+def _tile_impl(operand, *, reps):
+  out_shape = tuple(d * r for d, r in zip(operand.shape, reps))
+  bcast_shape = []
+  bcast_dims = []
+  for d, r in zip(operand.shape, reps):
+    if d == 1 or r == 1:
+      bcast_dims.append(len(bcast_shape))
+      bcast_shape.append(d * r)
+    else:
+      bcast_dims.append(len(bcast_shape) + 1)
+      bcast_shape.extend((r, d))
+  bcast = broadcast_in_dim(operand, tuple(bcast_shape), tuple(bcast_dims))
+  return reshape(bcast, out_shape)
+
+
+def _tile_transpose(ct, operand, *, reps):
+  assert ad.is_undefined_primal(operand)
+  if type(ct) is ad_util.Zero:
+    return ad_util.Zero(operand.aval)
+  reshape_shape = []
+  reduce_dims = []
+  for d, r in zip(operand.aval.shape, reps):
+    if r == 1:
+      reshape_shape.append(d)
+    elif d == 1:
+      reduce_dims.append(len(reshape_shape))
+      reshape_shape.append(r)
+    else:
+      reduce_dims.append(len(reshape_shape))
+      reshape_shape.extend((r, d))
+  reshaped_ct = reshape(ct, tuple(reshape_shape))
+  return [reduce_sum(reshaped_ct, tuple(reduce_dims))]
+
+
+def _tile_batching_rule(batched_args, batch_dims, *, reps):
+  (operand,) = batched_args
+  (bdim,) = batch_dims
+  if bdim is None:
+    return tile(operand, reps), None
+  reps = list(reps)
+  reps.insert(bdim, 1)
+  return tile(operand, reps), bdim
+
+
+tile_p = core.Primitive('tile')
+tile_p.def_impl(_tile_impl)
+tile_p.def_abstract_eval(_tile_abstract_eval)
+ad.deflinear2(tile_p, _tile_transpose)
+batching.primitive_batchers[tile_p] = _tile_batching_rule
+mlir.register_lowering(
+    tile_p, mlir.lower_fun(_tile_impl, multiple_results=False))
+
+
 def _clamp_shape_rule(min, operand, max):
   if min.shape and min.shape != operand.shape:
     raise TypeError("clamp requires min.shape == operand.shape or min.shape == "

jax/_src/numpy/lax_numpy.py

@@ -4519,11 +4519,13 @@ def tile(A: ArrayLike, reps: DimSize | Sequence[DimSize]) -> Array:
     reps_tup = tuple(reps)  # type: ignore[arg-type]
   reps_tup = tuple(operator.index(rep) if core.is_constant_dim(rep) else rep
                    for rep in reps_tup)
-  A_shape = (1,) * (len(reps_tup) - np.ndim(A)) + np.shape(A)
-  reps_tup = (1,) * (len(A_shape) - len(reps_tup)) + reps_tup
-  result = broadcast_to(reshape(A, [j for i in A_shape for j in [1, i]]),
-                        [k for pair in zip(reps_tup, A_shape) for k in pair])
-  return reshape(result, tuple(np.multiply(A_shape, reps_tup)))
+
+  # Prepend 1s to reps to match A.ndim
+  if len(reps_tup) < A.ndim:
+    reps_tup = (1,) * (A.ndim - len(reps_tup)) + reps_tup
+  if len(reps_tup) > A.ndim:
+    A = lax.expand_dims(A, list(range(len(reps_tup) - A.ndim)))
+  return lax.tile(A, reps_tup)
 
 def _concatenate_array(arr: ArrayLike, axis: int | None,
                        dtype: DTypeLike | None = None) -> Array:

jax/lax/__init__.py

@@ -229,6 +229,8 @@
   tan_p as tan_p,
   tanh as tanh,
   tanh_p as tanh_p,
+  tile as tile,
+  tile_p as tile_p,
   top_k as top_k,
   top_k_p as top_k_p,
   transpose as transpose,

tests/lax_autodiff_test.py

@@ -1180,5 +1180,21 @@ def testPowShapeMismatch(self):
     self.assertArraysEqual(actual, expected)
 
 
+  @jtu.sample_product(
+      [
+          dict(arg_shape=arg_shape, reps=reps)
+          for arg_shape, reps in [
+              [(3,), (2,)],
+              [(2, 3), (1, 2)],
+          ]
+      ],
+      dtype=grad_float_dtypes,
+  )
+  def testTileAutodiff(self, arg_shape, reps, dtype):
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: [rng(arg_shape, dtype)]
+    op = lambda x: lax.tile(x, reps)
+    check_grads(op, args_maker(), order=3, modes=["fwd", "rev"], eps=1.)
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())

tests/lax_test.py

@@ -1526,6 +1526,26 @@ def testBroadcastInDimAgainstNumpy(self, inshape, dtype, outshape, dimensions):
     numpy_op = lambda x: lax_reference.broadcast_in_dim(x, outshape, dimensions)
     self._CheckAgainstNumpy(numpy_op, op, args_maker)
 
+  @jtu.sample_product(
+      [
+          dict(arg_shape=arg_shape, reps=reps)
+          for arg_shape, reps in [
+              [(3,), (2,)],
+              [(2, 3), (1, 2)],
+              [(2, 3), (2, 1)],
+              [(2, 1, 3), (1, 2, 3)],
+          ]
+      ],
+      dtype=lax_test_util.default_dtypes,
+  )
+  def testTile(self, arg_shape, reps, dtype):
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: [rng(arg_shape, dtype)]
+    op = lambda x: lax.tile(x, reps)
+    numpy_op = lambda x: np.tile(x, reps)
+    self._CompileAndCheck(op, args_maker)
+    self._CheckAgainstNumpy(numpy_op, op, args_maker)
+
   @parameterized.parameters(
     {"inshape": inshape, "dimensions": dimensions, "error_type": error_type,
      "err_msg": err_msg}

tests/lax_vmap_test.py

@@ -791,5 +791,45 @@ def g(a, b):
     self.assertAllClose(output, expected, check_dtypes=False)
 
 
+  @jtu.sample_product(
+      [
+          dict(arg_shape=arg_shape, reps=reps)
+          for arg_shape, reps in [
+              [(3,), (2,)],
+              [(2, 3), (1, 2)],
+              [(2, 3), (2, 1)],
+              [(2, 1, 3), (1, 2, 3)],
+          ]
+      ],
+      in_axes=[0, 1, -1],
+      out_axes=[0, 1, -1],
+  )
+  def testTileBatching(self, arg_shape, reps, in_axes, out_axes):
+    rng = jtu.rand_default(self.rng())
+    dtype = np.float32
+    args_maker = lambda: [rng(arg_shape, dtype)]
+    op = lambda x: lax.tile(x, reps)
+    args = args_maker()
+
+    # Construct batched arguments based on in_axes
+    if in_axes == 0:
+      batched_args = [jnp.stack([arg, arg], axis=0) for arg in args]
+    elif in_axes == 1:
+      batched_args = [jnp.stack([arg, arg], axis=1) for arg in args]
+    else: # in_axes == -1
+      batched_args = [jnp.stack([arg, arg], axis=-1) for arg in args]
+
+    # Compute expected output
+    out = op(*args)
+    if out_axes == 0:
+      expected = jnp.stack([out, out], axis=0)
+    elif out_axes == 1:
+      expected = jnp.stack([out, out], axis=1)
+    else: # out_axes == -1
+      expected = jnp.stack([out, out], axis=-1)
+
+    actual = jax.vmap(op, in_axes=in_axes, out_axes=out_axes)(*batched_args)
+    self.assertAllClose(expected, actual)
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())