Repository: google/jax

Issue:
consider duck typing for jnp.ones_like, other similar functions
right now, `jnp.ones_like(jax.ShapeDtypeStruct(shape=(1, 2, 3), dtype=jnp.float32)` throws a type error
it probably shouldn't, and in fact `jnp.ones_like` should probably work on anything that has `shape`, `dtype`, etc attributes


Discussion:
`jnp.ones_like` et al would ideally get shape, dtype, etc from `x` whatever `x` is and only check that the shape, dtype, etc are themselves well typed


Seems reasonable! (In fact I vaguely recall this coming up before...)

What do you think, @jakevdp ?
Sure, sounds reasonable to me. We have to be somewhat careful because some valid inputs won't have a `dtype` or `shape` defined (e.g. Python scalars) and some non-array inputs we want to be errors (e.g. lists and tuples).
Maybe we can try to look for `shape` and `dtype` attributes first, and if that fails fall back to the current path.
I just ran into a case where this would be nice as well. Would it be possible for me to make a PR for this or are there non-obvious things that could get broken?

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -68,7 +68,7 @@
 from jax._src.lib.mlir.dialects import chlo
 from jax._src.lib.mlir.dialects import hlo
 from jax._src.sharding_impls import PmapSharding
-from jax._src.typing import Array, ArrayLike, DTypeLike, Shape
+from jax._src.typing import Array, ArrayLike, DuckTypedArray, DTypeLike, Shape
 from jax._src.util import (cache, safe_zip, safe_map, canonicalize_axis,
                            split_list)
 
@@ -1330,7 +1330,8 @@ def expand_dims(array: ArrayLike, dimensions: Sequence[int]) -> Array:
 
 ### convenience wrappers around traceables
 
-def full_like(x: ArrayLike, fill_value: ArrayLike, dtype: Optional[DTypeLike] = None,
+def full_like(x: Union[ArrayLike, DuckTypedArray],
+              fill_value: ArrayLike, dtype: Optional[DTypeLike] = None,
               shape: Optional[Shape] = None) -> Array:
   """Create a full array like np.full based on the example array `x`.
 
@@ -1344,7 +1345,7 @@ def full_like(x: ArrayLike, fill_value: ArrayLike, dtype: Optional[DTypeLike] =
     An ndarray with the same shape as `x` with its entries set equal to
     `fill_value`, similar to the output of np.full.
   """
-  fill_shape = np.shape(x) if shape is None else canonicalize_shape(shape)
+  fill_shape = np.shape(x) if shape is None else canonicalize_shape(shape)  # type: ignore[arg-type]
   weak_type = dtype is None and dtypes.is_weakly_typed(x)
   dtype = dtype or _dtype(x)
   if dtypes.is_opaque_dtype(dtype):
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -60,7 +60,7 @@
 from jax._src.numpy import ufuncs
 from jax._src.numpy import util
 from jax._src.numpy.vectorize import vectorize
-from jax._src.typing import Array, ArrayLike, DimSize, DType, DTypeLike, Shape
+from jax._src.typing import Array, ArrayLike, DimSize, DuckTypedArray, DType, DTypeLike, Shape
 from jax._src.util import (unzip2, subvals, safe_zip,
                            ceil_of_ratio, partition_list,
                            canonicalize_axis as _canonicalize_axis)
@@ -320,7 +320,7 @@ def isscalar(element: Any) -> bool:
 iterable = np.iterable
 
 @util._wraps(np.result_type)
-def result_type(*args: ArrayLike) -> DType:
+def result_type(*args: Any) -> DType:
   return dtypes.result_type(*args)
 
 
@@ -2077,9 +2077,11 @@ def copy(a: ArrayLike, order: Optional[str] = None) -> Array:
 
 
 @util._wraps(np.zeros_like)
-def zeros_like(a: ArrayLike, dtype: Optional[DTypeLike] = None,
+def zeros_like(a: Union[ArrayLike, DuckTypedArray],
+               dtype: Optional[DTypeLike] = None,
                shape: Any = None) -> Array:
-  util.check_arraylike("zeros_like", a)
+  if not (hasattr(a, 'dtype') and hasattr(a, 'shape')):  # support duck typing
+    util.check_arraylike("ones_like", a)
   dtypes.check_user_dtype_supported(dtype, "zeros_like")
   if shape is not None:
     shape = canonicalize_shape(shape)
@@ -2087,9 +2089,11 @@ def zeros_like(a: ArrayLike, dtype: Optional[DTypeLike] = None,
 
 
 @util._wraps(np.ones_like)
-def ones_like(a: ArrayLike, dtype: Optional[DTypeLike] = None,
+def ones_like(a: Union[ArrayLike, DuckTypedArray],
+              dtype: Optional[DTypeLike] = None,
               shape: Any = None) -> Array:
-  util.check_arraylike("ones_like", a)
+  if not (hasattr(a, 'dtype') and hasattr(a, 'shape')):  # support duck typing
+    util.check_arraylike("ones_like", a)
   dtypes.check_user_dtype_supported(dtype, "ones_like")
   if shape is not None:
     shape = canonicalize_shape(shape)
@@ -2099,9 +2103,11 @@ def ones_like(a: ArrayLike, dtype: Optional[DTypeLike] = None,
 @util._wraps(np.empty_like, lax_description="""\
 Because XLA cannot create uninitialized arrays, the JAX version will
 return an array initialized with zeros.""")
-def empty_like(prototype: ArrayLike, dtype: Optional[DTypeLike] = None,
+def empty_like(prototype: Union[ArrayLike, DuckTypedArray],
+               dtype: Optional[DTypeLike] = None,
                shape: Any = None) -> Array:
-  util.check_arraylike("empty_like", prototype)
+  if not (hasattr(prototype, 'dtype') and hasattr(prototype, 'shape')):  # support duck typing
+    util.check_arraylike("ones_like", prototype)
   dtypes.check_user_dtype_supported(dtype, "empty_like")
   return zeros_like(prototype, dtype=dtype, shape=shape)
 
@@ -2119,17 +2125,21 @@ def full(shape: Any, fill_value: ArrayLike,
 
 
 @util._wraps(np.full_like)
-def full_like(a: ArrayLike, fill_value: ArrayLike, dtype: Optional[DTypeLike] = None,
+def full_like(a: Union[ArrayLike, DuckTypedArray],
+              fill_value: ArrayLike, dtype: Optional[DTypeLike] = None,
               shape: Any = None) -> Array:
+  if hasattr(a, 'dtype') and hasattr(a, 'shape'):  # support duck typing
+    util.check_arraylike("full_like", 0, fill_value)
+  else:
+    util.check_arraylike("full_like", a, fill_value)
   dtypes.check_user_dtype_supported(dtype, "full_like")
-  util.check_arraylike("full_like", a, fill_value)
   if shape is not None:
     shape = canonicalize_shape(shape)
   if ndim(fill_value) == 0:
     return lax.full_like(a, fill_value, dtype, shape)
   else:
-    shape = np.shape(a) if shape is None else shape
-    dtype = result_type(a) if dtype is None else dtype
+    shape = np.shape(a) if shape is None else shape  # type: ignore[arg-type]
+    dtype = result_type(a) if dtype is None else dtype  # type: ignore[arg-type]
     return broadcast_to(asarray(fill_value, dtype=dtype), shape)
 
 
diff --git a/jax/_src/typing.py b/jax/_src/typing.py
--- a/jax/_src/typing.py
+++ b/jax/_src/typing.py
@@ -57,6 +57,12 @@ def dtype(self) -> DType: ...
 DimSize = Union[int, Any]  # extensible
 Shape = Sequence[DimSize]
 
+class DuckTypedArray(Protocol):
+  @property
+  def dtype(self) -> DType: ...
+  @property
+  def shape(self) -> Shape: ...
+
 # Array is a type annotation for standard JAX arrays and tracers produced by
 # core functions in jax.lax and jax.numpy; it is not meant to include
 # future non-standard array types like KeyArray and BInt. It is imported above.


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -2621,6 +2621,13 @@ def testZerosOnesLike(self, func, shape, in_dtype, out_shape, out_dtype):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  def testDuckTypedLike(self):
+    x = jax.ShapeDtypeStruct((1, 2, 3), np.dtype("int32"))
+    self.assertArraysEqual(jnp.zeros_like(x), jnp.zeros(x.shape, x.dtype))
+    self.assertArraysEqual(jnp.ones_like(x), jnp.ones(x.shape, x.dtype))
+    self.assertArraysEqual(jnp.empty_like(x), jnp.empty(x.shape, x.dtype))
+    self.assertArraysEqual(jnp.full_like(x, 2), jnp.full(x.shape, 2, x.dtype))
+
   @jtu.sample_product(
     [dict(func=func, args=args)
      for func, args in [("full_like", (-100,)), ("ones_like", ()), ("zeros_like", ())]