Latest openpilot .onnx does not compile

[haraschax]

FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python compile2.py https://github.com/haraschax/filedump/raw/master/notinygradcompile.onnx

Produces:

batman@workstation-harald:~/xx/openpilot/tinygrad_repo/openpilot$ FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python compile2.py https://github.com/haraschax/filedump/raw/master/notinygradcompile.onnx
https://github.com/haraschax/filedump/raw/master/notinygradcompile.onnx: 100%|█| 50.7M/50.7M [00:00<00:
Traceback (most recent call last):
  File "/home/batman/xx/openpilot/tinygrad_repo/openpilot/compile2.py", line 104, in <module>
    schedule, schedule_independent, inputs = get_schedule(onnx_data)
                                             ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/openpilot/compile2.py", line 33, in get_schedule
    ret: Tensor = next(iter(run_onnx(inputs).values())).cast(dtypes.float32).contiguous()
                            ^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/extra/onnx.py", line 202, in run_onnx
    ret = real_fxn(*inp, **opt)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/extra/onnx_ops.py", line 84, in Reshape
    return data.reshape([int(x) if x != 0 else (0 if allowzero else data.shape[i]) for i,x in enumerate(safe_numpy(shape))])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/tinygrad/tensor.py", line 669, in reshape
    return F.Reshape.apply(self, shape=new_shape) if new_shape != self.shape else self
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/tinygrad/tensor.py", line 37, in apply
    ret.lazydata, ret.requires_grad, ret.grad = ctx.forward(*[t.lazydata for t in x], **kwargs), ctx.requires_grad, None
                                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/tinygrad/function.py", line 179, in forward
    return x.reshape(shape)
           ^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/tinygrad/lazy.py", line 208, in reshape
    def reshape(self, arg:Tuple[sint, ...]): return self._view(self.st.reshape(arg))
                                                               ^^^^^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/tinygrad/shape/shapetracker.py", line 114, in reshape
    if getenv("MERGE_VIEW", 1) and (new_view := self.views[-1].reshape(new_shape)) is not None: return ShapeTracker(self.views[0:-1] + (new_view,))
                                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/tinygrad/shape/view.py", line 271, in reshape
    raise ValueError(f"size mismatched, can't reshape {self.shape=} -> {new_shape=}")
ValueError: size mismatched, can't reshape self.shape=(1, 1) -> new_shape=(1, 1, 2)

[geohot]

Reproduced issue.

@geohotstan want to fix?

Running with DEBUGONNX=1 gets:

...
611: op Mul shape [(2,), ()] opt {}
612: op Equal shape [(2,), (2,)] opt {}
613: op Where shape [(2,), (2,), (2,)] opt {}
614: op Expand shape [(1, 1), (2,)] opt {}
615: op Constant shape [] opt {'value': <Tensor <LB GPU (1,) long (<LoadOps.COPY: 3>, None)> on GPU with grad None>}
616: op Unsqueeze shape [(1, 1), (1,)] opt {}
617: op Constant shape [] opt {'value': <Tensor <LB GPU (1,) long (<LoadOps.COPY: 3>, None)> on GPU with grad None>}
618: op ConstantOfShape shape [(1,)] opt {'value': <Tensor <LB GPU (1,) long (<LoadOps.COPY: 3>, None)> on GPU with grad None>}
619: op Constant shape [] opt {'value': <Tensor <LB GPU () long ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),))> on GPU with grad None>}
620: op Mul shape [(2,), ()] opt {}
621: op Equal shape [(2,), (2,)] opt {}
622: op Where shape [(2,), (2,), (2,)] opt {}
623: op Expand shape [(1,), (2,)] opt {}
624: op Constant shape [] opt {'value': <Tensor <LB GPU (1,) long (<LoadOps.COPY: 3>, None)> on GPU with grad None>}
625: op Unsqueeze shape [(1, 1), (1,)] opt {}
626: op Concat shape [(1, 1, 1), (1, 1, 1)] opt {'axis': -1}
627: op Shape shape [(1, 2)] opt {}
628: op Constant shape [] opt {'value': <Tensor <LB GPU (1,) long (<LoadOps.COPY: 3>, None)> on GPU with grad None>}
629: op Constant shape [] opt {'value': <Tensor <LB GPU (1,) long (<LoadOps.COPY: 3>, None)> on GPU with grad None>}
630: op Constant shape [] opt {'value': <Tensor <LB GPU (1,) long (<LoadOps.COPY: 3>, None)> on GPU with grad None>}
631: op Slice shape [(2,), (1,), (1,), (1,)] opt {}
632: op Concat shape [(2,), (1,)] opt {'axis': 0}
633: op Reshape shape [(1, 1), (3,)] opt {'allowzero': 0}

[geohot]

As discussed in the first fix, we now need ScatterND to make this work.

[haraschax]

It fails cleanly now:

(.venv) batman@workstation-harald:~/xx/openpilot/tinygrad_repo/openpilot$ FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python compile2.py https://github.com/haraschax/filedump/raw/master/notinygradcompile.onnx
UNSUPPORTED ScatterND ['/action_block/action_block_out/Gemm_output_0', '/action_block/Concat_output_0', '/action_block/Reshape_output_0'] ['/action_block/ScatterND_output_0']
Traceback (most recent call last):
  File "/home/batman/xx/openpilot/tinygrad_repo/openpilot/compile2.py", line 104, in <module>
    schedule, schedule_independent, inputs = get_schedule(onnx_data)
                                             ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/openpilot/compile2.py", line 33, in get_schedule
    ret: Tensor = next(iter(run_onnx(inputs).values())).cast(dtypes.float32).contiguous()
                            ^^^^^^^^^^^^^^^^
  File "/home/batman/xx/openpilot/tinygrad_repo/extra/onnx.py", line 205, in run_onnx
    raise Exception(f"op_type {n.op_type} not supported")
Exception: op_type ScatterND not supported

Feel free to close if tinygrad does not want to support ScatterND. I can work around that for this model.

[geohotstan]

tinygrad currently doesn't support fancy indexing setitem, which would be the best way of implementing ScatterND.
I've been looking into supporting full setitem and I do think it's possible. It just might take some time to figure out.
So if supporting this is urgent, then maybe the work around would be best. 😄