Add support for pytorch.tensor

[HernandoR]

I'm interested in adding support for pytorch tensor, such that users can use it directly without constantly convert the data type.

The coding is trivial, I'm already taking hands on it. but this may introduce optional dependencies of pytorch. I'm not sure how to edit the setup.py file.

Please let me know what you think of this feature and if you can help me with the dependency problem

[AlexanderFabisch]

Hi @HernandoR,

I personally like JAX more than PyTorch, but I found it quite complicated to support JAX and NumPy in one library. That's why I created jaxtransform3d: https://github.com/AlexanderFabisch/jaxtransform3d . Just in case that switching the framework is an option.

Only supporting PyTorch is a bit specific. I thought about using the array API (which is implemented by PyTorch and several other libraries), but I found it really hard to implement while keeping the readability of the code.

Do you have an example of how you would modify the code to support PyTorch? I think one of the main problems is creating new tensors/arrays, which often needs a library-specific function call.

[AlexanderFabisch]

A comment on the additional dependency: it must be an optional dependency. I want to limit the dependencies as much as possible.

[HernandoR]

Similar to open3d, I would like to use a pytransform3d.t.* namespace, within which all api would be same (hopfully) except implemented by pytorch and tensors. Such that the users could simply change the import area to specify if they want numpy or torch

[HernandoR]

As for array api, I think pytorch is not yet fully compatible with the array api: pytorch array api tracer; I could try to implement base on array api if we could conclude that the current support is good enough.

[AlexanderFabisch]

This is a huge project that would add considerable maintenance overhead to the library. I am not entirely sure about it.

Would it be possible to start this in a separate repository? We could still make it available under pytransform3d.t.* if it supports the same interface.

The better option would be array API support in my opinion though. I would expect the compatibility issue to be solved in the future.

Another idea would be to limit the whole project to batch_rotations or trajectories for now. It is a smaller set of functions and they are all vectorized. I'm pretty sure that's good when you work with GPU support.

[HernandoR]

@AlexanderFabisch

I adopted a module to array_api. During which I encountered 2 problems.

1. Most of the signatures accept array-like which includes list as input. But array_api cannot determine a namespace out of it.
2. Some basic functions accept float or array and return float or array respectively. array_api must be used in full array pipeline.

For both problems, I think we can assume a fallback namespace such as numpy, i.e., we could always transfer the input to a numpy array if the user passed a list or float.

Does this assumption make sense in our use case?

[AlexanderFabisch]

Most of the signatures accept array-like which includes list as input. But array_api cannot determine a namespace out of it.

Usually you would cast lists to numpy arrays with np.asarray. I don't know what the standard approach to do this is. Maybe have a look at scipy or scikit-learn to see how they are doing it.

Some basic functions accept float or array and return float or array respectively. array_api must be used in full array pipeline.

Do you have examples? Maybe not every function has to support the array API. I can only think of matrix_from_angle, in which I wouldn't want to cast to a numpy array for performance reasons.

[HernandoR]

An example:

def norm_angle(a):
    """Normalize angle to (-pi, pi].

    It is worth noting that using `numpy.ceil` to normalize angles will lose
    more digits of precision as angles going larger but can keep more digits
    of precision when angles are around zero. In common use cases, for example,
    -10.0*pi to 10.0*pi, it performs well.

    For more discussions on numerical precision:
    https://github.com/dfki-ric/pytransform3d/pull/263

    Parameters
    ----------
    a : float or array-like, shape (n,)
        Angle(s) in radians

    Returns
    -------
    a_norm : float or array, shape (n,)
        Normalized angle(s) in radians
    """
    a = np.asarray(a, dtype=np.float64)
    return a - (np.ceil((a + np.pi) / two_pi) - 1.0) * two_pi

Current behaviour is to return a float or an array. To support array api and maintain compatibility, we would make an assumption use numpy's asarray if we could not infer a namespace from the input.

As to my understanding of the array api (please point out if I misunderstood), the array api suggest substractions return an array. So I'm not sure what is expected in array api , i.e.
assert type(norm_angle(1.0))==xp.float64 or assert type(norm_angle(1.0))==xp.array

[AlexanderFabisch]

The np.asarray call is already in the function? If so, then I don't see an issue with this particular function as it returns a 0d array, not a float, right?

[HernandoR]

My paste here is the unmodified version, and it returns a np.float64 object. M here is that, the current implementation uses np.asarray specifically. If we would like to support array api in this function, we have to infer the namespace from the input. i.e.

def norm_angle(a):
    xp= xarray.array_namespace(a)
    a=xp.asarray(a, dtype=np.float64)
    two_pi=2.0*xp.pi
    return a - (xp.ceil((a + xp.pi) / two_pi) - 1.0) * two_pi

but this implement could not deal with float input (xp.float64 is fine though), so we end with something like

def norm_angle(a):
    if isinstance(a, float):
        a=np.asarray(a, dtype=np.float64)
    xp=xarray.array_namespace(a)
    a=xp.asarray(a, dtype=np.float64)
    two_pi=2.0*xp.pi
    return a - (xp.ceil((a + xp.pi) / two_pi) - 1.0) * two_pi

which implies we will return a np.float64 for float input rather than any other namespace the user may would like to use.

[AlexanderFabisch]

If the user wants a specific array type as output, it is perfectly fine to expect them to pass an argument of that type. I don't see any problem with that.