Reverts f6101c5

PiperOrigin-RevId: 825412424

Author: Google-ML-Automation

jax/_src/lax/lax.py

@@ -6729,101 +6729,6 @@ 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,13 +4519,11 @@ 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)
-
-  # 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)
+  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)))
 
 def _concatenate_array(arr: ArrayLike, axis: int | None,
                        dtype: DTypeLike | None = None) -> Array:

jax/lax/__init__.py

@@ -229,8 +229,6 @@
   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,21 +1180,5 @@ 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,26 +1526,6 @@ 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,45 +791,5 @@ 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())