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.

[jonaspleyer]

The traits utilized in approx are currently only able to produce bools. If you are interested in providing methods which yield Array-like types, we would have to extend the traits. This would require a redesign but could be doable. In particular, other crates might be interested in such functionality as well.

Possible Design:

trait AbsDiffEq<Rhs=Self, Ret=bool>: PartialEq<Rhs>
{
    type Epsilon;

    fn default_epsilon() -> Self::Epsilon;
    fn abs_diff_eq(&self, other: &Rhs, epsilon: Self::Epsilon) -> Ret;
    ..
}

Then we could have implementations of the following forms (sketch)

impl<F, D> AbsDiffEq for ArrayRef<F, D> {
    type Epsilon = F;
    fn abs_diff_eq(&self, other: &Rhs, epsilon: Self::Epsilon) -> bool {..} // <-- Returns bool
    ..
}

impl<F, D> AbsDiffEq<Ret=Array<F, D>> for ArrayRef<F, D> {
    type Epsilon = F;
    fn abs_diff_eq(&self, other: &Rhs, epsilon: Self::Epsilon) -> Array<F, D> {..} // <-- Returns Array
    ..
}

So far, this seems doable but I have not fully thought through all the possible consequences that this might have. It would also be important to note that implementations of the following form

impl<F> AbsDiffEq<Rhs=Array2<F>, Ret=Array2<F>> for Array3<F> {..}

could be realized in many different ways and lead to confusion in practice. I am not sure of how a particular type signature might improve the clarity of such comparisons or only introduce unwieldy type-juggling discomfort.