An `Unknown` type is returned from `array-api-extra`

[paddyroddy]

Add your issue here

If the # ty: ignore[invalid-return-type] comment is removed from the following code

glass/glass/shells.py

Lines 520 to 529 in 7e620f7

	zr = xpx.union1d(w.za, z_)
	fr = glass.arraytools.ndinterp(
	zr,
	z,
	f,
	left=0.0,
	right=0.0,
	) * glass.arraytools.ndinterp(zr, w.za, w.wa)
	return zr, fr # ty: ignore[invalid-return-type]

then we get this error upon running ty check

error[invalid-return-type]: Return type does not match returned value
   --> glass/shells.py:529:12
    |
529 |     return zr, fr
    |            ^^^^^^ expected `tuple[ndarray[tuple[Any, ...], dtype[float64]] | jax._src.basearray.Array | array_api_strict._array_object.Array, ndarray[tuple[Any, ...], dtype[float64]] | jax._src.basearray.Array | array_api_strict._array_object.Array]`, found `tuple[array_api_extra._lib._utils._typing.Array, Unknown]`
    |
   ::: glass/shells.py:482:6
    |
482 | ) -> tuple[FloatArray, FloatArray]:
    |      ----------------------------- Expected `tuple[ndarray[tuple[Any, ...], dtype[float64]] | jax._src.basearray.Array | array_api_strict._array_object.Array, ndarray[tuple[Any, ...], dtype[float64]] | jax._src.basearray.Array | array_api_strict._array_object.Array]` because of return type
    |
info: the first tuple element is not compatible: `Array` is not assignable to `ndarray[tuple[Any, ...], dtype[float64]] | jax._src.basearray.Array | array_api_strict._array_object.Array`

The Unknown type is coming from the xpx.union1d call. This could be "fixed" by changing line 520 to zr = xp.asarray(xpx.union1d(w.za, z_)).

[lucascolley]

The Unknown type is coming from the xpx.union1d call. This could be "fixed" by changing line 520 to zr = xp.asarray(xpx.union1d(w.za, z_)).

Are you sure? From that error message, it looks to me like zr is of type array_api_extra._lib._utils._typing.Array and it is fr (coming from glass.arraytools.ndinterp) that is of type Unknown.

[paddyroddy]

So I've installed the ty extension

so yes the fr is coming out as unknown. ndinterp is using uxpx.apply_along_axis

glass/glass/arraytools.py

Lines 143 to 151 in 682b8e6

	return uxpx.apply_along_axis(
	uxpx.interp,
	(x, xq),
	axis,
	fq,
	left=left,
	right=right,
	period=period,
	)

. The output of uxpx.apply_along_axis should be an array

glass/glass/_array_api_utils.py

Lines 385 to 398 in 682b8e6

	xp = array_api_compat.array_namespace(arr, *func_inputs, use_compat=False)
	if xp.__name__ in {"numpy", "jax.numpy"}:
	func1d = functools.partial(func, *func_inputs)
	return xp.apply_along_axis(func1d, axis, arr, *args, **kwargs)
	# If any other backend use default
	dxp = default_xp(xp.__name__)
	inputs_dxp = (dxp.asarray(inp) for inp in func_inputs)
	func1d = functools.partial(func, *inputs_dxp)
	return xp.asarray(
	dxp.apply_along_axis(func1d, axis, arr, *args, **kwargs),
	copy=True,
	)

.

[lucascolley]

is it due to the * operator being undefined between some of the elements of the AnyArray union?

[paddyroddy]

Ah, yes it could well be that. Will have to see if I can type that somehow.

[lucascolley]

I would be tempted to just use a cast to AnyArray for now, pending array-api-typing