Update src/array_api_extra/_delegation.py

NeilGirdhar · crusaderky · web-flow · commit e883014fbb61 · 2025-06-02T08:22:13.000-04:00
Co-authored-by: Guido Imperiale &lt;crusaderky@gmail.com&gt;
diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
@@ -175,7 +175,6 @@ def one_hot(
         ]
     # Delegate where possible.
     if is_jax_namespace(xp):
-        assert is_jax_array(x)
         from jax.nn import one_hot as jax_one_hot
 
         return jax_one_hot(x, num_classes, dtype=dtype, axis=axis)

Original file line number	Diff line number	Diff line change
`@@ -175,7 +175,6 @@ def one_hot(`
`175`	`175`	`]`
`176`	`176`	`# Delegate where possible.`
`177`	`177`	`if is_jax_namespace(xp):`
`178`		`- assert is_jax_array(x)`
`179`	`178`	`from jax.nn import one_hot as jax_one_hot`
`180`	`179`
`181`	`180`	`return jax_one_hot(x, num_classes, dtype=dtype, axis=axis)`