best way to initialize jax ndarray from jnp.ndarray, np.ndarray or python float #10212
-
Hi all, I am implementing a function which accepts def fn(*, v: Union[jnp.ndarray, np.ndarray, list(float), float], ...):
v = jnp.asarray(v) Suppose I am using GPU backend. The input parameter I have read some discussion like #4486 and #5012 . On TF, it may be like this:
|
Beta Was this translation helpful? Give feedback.
Replies: 1 comment
-
I think |
Beta Was this translation helpful? Give feedback.
I think
jnp.asarray
is okay, and it is equivalent tojnp.array
now.https://github.com/google/jax/blob/5fdad0ebf5b6bd97df01eebe86d3ecb49f193d2a/jax/_src/numpy/lax_numpy.py#L1865-L1868
https://github.com/google/jax/blob/5fdad0ebf5b6bd97df01eebe86d3ecb49f193d2a/jax/_src/numpy/lax_numpy.py#L1819-L1833
BTW,
jnp.ndarray
is actually a abstract class.DeviceArray
is the concrete class.