OneHot broken?

[antoinediez]

Dear Keops developers,

Together with @amic-github, we have been trying to implement a reduction formula which involves a one_hot LazyTensor but we got a very strange error which makes me think that something has been broken in a recent update.

The following code

import torch
from pykeops.torch import LazyTensor

use_cuda = torch.cuda.is_available()
dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor

x=torch.tensor([0.,5.,10.]).type(dtype).reshape(3,1)
xi=LazyTensor(x[:,None,:])
y = xi.one_hot(42)
y.sum(1)

returns the following error on Google Colab

[KeOps] Compiling cuda jit compiler engine ... OK
[pyKeOps] Compiling nvrtc binder for python ... OK
OneHot

---------------------------------------------------------------------------

NameError                                 Traceback (most recent call last)

[<ipython-input-2-3a5c9c813431>](https://localhost:8080/#) in <module>()
      8 xi=LazyTensor(x[:,None,:])
      9 y = xi.one_hot(42)
---> 10 y.sum(1)

12 frames

[/usr/local/lib/python3.7/dist-packages/pykeops/common/lazy_tensor.py](https://localhost:8080/#) in sum(self, axis, dim, **kwargs)
   1798             return self.unary("Sum", dimres=1)
   1799         else:
-> 1800             return self.reduction("Sum", axis=axis, **kwargs)
   1801 
   1802     def sum_reduction(self, axis=None, dim=None, **kwargs):

[/usr/local/lib/python3.7/dist-packages/pykeops/common/lazy_tensor.py](https://localhost:8080/#) in reduction(self, reduction_op, other, opt_arg, axis, dim, call, is_complex, **kwargs)
    753             )
    754         if call and len(res.symbolic_variables) == 0 and res._dtype is not None:
--> 755             return res()
    756         else:
    757             return res

[/usr/local/lib/python3.7/dist-packages/pykeops/common/lazy_tensor.py](https://localhost:8080/#) in __call__(self, *args, **kwargs)
    935             args = (self.other.variables[0],)
    936 
--> 937         return self.callfun(*args, *self.variables, **self.kwargs)
    938 
    939     def __str__(self):

[/usr/local/lib/python3.7/dist-packages/pykeops/torch/generic/generic_red.py](https://localhost:8080/#) in __call__(self, backend, device_id, ranges, out, *args)
    634             ny,
    635             out,
--> 636             *args
    637         )
    638 

[/usr/local/lib/python3.7/dist-packages/pykeops/torch/generic/generic_red.py](https://localhost:8080/#) in forward(ctx, formula, aliases, backend, dtype, device_id_request, ranges, optional_flags, rec_multVar_highdim, nx, ny, out, *args)
     87             dtype,
     88             "torch",
---> 89             optional_flags,
     90         ).import_module()
     91 

[/usr/local/lib/python3.7/dist-packages/keopscore/utils/Cache.py](https://localhost:8080/#) in __call__(self, *args)
     66                     self.library[str_id] = self.cls(params, fast_init=True)
     67                 else:
---> 68                     obj = self.cls(*args)
     69                     self.library_params[str_id] = obj.params
     70                     self.library[str_id] = obj

[/usr/local/lib/python3.7/dist-packages/pykeops/common/keops_io/LoadKeOps_nvrtc.py](https://localhost:8080/#) in __init__(self, fast_init, *args)
     13 class LoadKeOps_nvrtc_class(LoadKeOps):
     14     def __init__(self, *args, fast_init=False):
---> 15         super().__init__(*args, fast_init=fast_init)
     16 
     17     def init_phase2(self):

[/usr/local/lib/python3.7/dist-packages/pykeops/common/keops_io/LoadKeOps.py](https://localhost:8080/#) in __init__(self, fast_init, *args)
     16             self.params = args[0]
     17         else:
---> 18             self.init(*args)
     19         self.dimout = self.params.dim
     20         self.tagIJ = self.params.tagI

[/usr/local/lib/python3.7/dist-packages/pykeops/common/keops_io/LoadKeOps.py](https://localhost:8080/#) in init(self, tagCPUGPU, tag1D2D, tagHostDevice, use_ranges, device_id_request, formula, aliases, nargs, dtype, lang, optional_flags)
    139             tag1D2D,
    140             self.params.use_half,
--> 141             device_id_request,
    142         )
    143 

[/usr/local/lib/python3.7/dist-packages/keopscore/utils/Cache.py](https://localhost:8080/#) in __call__(self, *args)
     25         str_id = "".join(list(str(arg) for arg in args)) + str(env_param)
     26         if not str_id in self.library:
---> 27             self.library[str_id] = self.fun(*args)
     28         return self.library[str_id]
     29 

[/usr/local/lib/python3.7/dist-packages/keopscore/get_keops_dll.py](https://localhost:8080/#) in get_keops_dll_impl(map_reduce_id, red_formula_string, enable_chunks, enable_finalchunks, mul_var_highdim, aliases, *args)
     91         set_enable_finalchunk(enable_finalchunks)
     92         set_mult_var_highdim(mul_var_highdim)
---> 93         red_formula = GetReduction(red_formula_string, aliases)
     94         if use_final_chunks(red_formula) and map_reduce_id != "GpuReduc2D":
     95             use_chunk_mode = 2

[/usr/local/lib/python3.7/dist-packages/keopscore/formulas/GetReduction.py](https://localhost:8080/#) in __new__(self, red_formula_string, aliases)
     25                     varname, var = alias.split("=")
     26                     aliases_dict[varname] = eval(var)
---> 27             reduction = eval(red_formula_string, globals(), aliases_dict)
     28             GetReduction.library[string_id_hash] = reduction
     29             return reduction

/usr/local/lib/python3.7/dist-packages/keopscore/formulas/GetReduction.py in <module>()

NameError: name 'OneHot' is not defined

Note also that the answer of @jeanfeydy in #24 (comment) which involves the one_hot function returns the same error (but I guess it did work back in 2019).

What do you think?

Thanks!

[jeanfeydy]

Hi @antoinediez,
Thanks for the catch - it was a very simple typo on our side, we just forgot to include OneHot in the list of supported operations when rewriting the engine from KeOps v1 to KeOps v2...
Please let me know if you encounter any other bug / unexpected behaviour, this is always very appreciated :-)
See you soon!
Jean