Array comparison operators/functions

[Qazalbash]

Hi,

I come from a Python background. I have used numpy for array-based computing and recently started working on JAX. They have very easy-to-use array comparison functions ¹²³ to check the relative, partial order etc. I find a lack of these functions/operators here. It would be a great enhancement and would help in testing code without writing a bunch of lines to compare errors or check the relationship between two arrays.

Until now I have been doing things like this,

assert!(distance.iter().all(|&d| d >= -ERR_MARGIN));

With the feature, it should more or less look like this,

assert!(distance.geq(-ERR_MARGIN));

Footnotes

1. jax.numpy.allclose ↩

2. jax.numpy.isclose ↩

3. jax.numpy.equal ↩

[nilgoyette]

You might already know, but enabling the approx feature will add some methods. They won't be named as you want though. I can find abs_diff_eq, relative_eq in the documentation, but there are also the assert versions: assert_relative_eq, assert_abs_diff_eq.

As for the element-wise method, I doubt that the maintainers want to add specific method like those, especially when they are so short and simple.

[akern40]

I'd have to check a little more closely, but I think we only implement these for array-to-array comparisons. Maybe we should consider implications for array-to-scalar, along with the approx traits for those?

[Qazalbash]

@akern40 along with array-to-scalar comparison, there should be array-to-array comparison in different axes too.