From da5643d02410fee86cb44d32bf3c5333407f14c0 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Fri, 24 Feb 2023 12:52:23 -0800
Subject: [PATCH] rest of tests shouid be made to pass

---
 extra/onnx_ops.py                  |  4 +++
 test/external_test_onnx_backend.py | 42 +++++++++++++++++++++++-------
 tinygrad/tensor.py                 |  2 +-
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/extra/onnx_ops.py b/extra/onnx_ops.py
index cf8840f63e..ae62068b01 100644
--- a/extra/onnx_ops.py
+++ b/extra/onnx_ops.py
@@ -65,6 +65,7 @@ def Dropout(data, ratio=0.5, training_mode=False, seed=None):
   return data * mask * (1/(1.0 - ratio)), mask
 
 def Shape(data, end=None, start=0): return list(data.shape)[start:end]
+def Size(data): return prod(data.shape)
 
 # TODO: this doesn't match Tensor.flatten behavior
 def Flatten(input, axis=1):
@@ -124,3 +125,6 @@ def Tile(input, repeats):
   expand_shape = [x for r,s in zip(repeats_, input.shape) for x in [r,s]]
   final_shape = [r*s for r,s in zip(repeats_, input.shape)]
   return input.reshape(new_shape).expand(expand_shape).reshape(final_shape)
+
+def Range(start, limit, delta): return Tensor.arange(safe_numpy(limit)[0], safe_numpy(start)[0], safe_numpy(delta)[0])
+def Where(condition, X, Y): return condition*X + (1-condition)*Y
diff --git a/test/external_test_onnx_backend.py b/test/external_test_onnx_backend.py
index 2ab876f7bd..75032086b4 100644
--- a/test/external_test_onnx_backend.py
+++ b/test/external_test_onnx_backend.py
@@ -44,18 +44,37 @@ backend_test.exclude('test_max_*')
 # add support for SoftmaxCrossEntropyLoss and NegativeLogLikelihoodLoss
 backend_test.exclude('test_sce_*')
 
-# no optimizers (add them)
+# no support for reduce with multiply (needs llop)
+backend_test.exclude('test_reduce_prod_*')
+
+# no optimizers (add them?)
 backend_test.exclude('test_adagrad_*')
 backend_test.exclude('test_adam_*')
 backend_test.exclude('test_nesterov_momentum_*')
+backend_test.exclude('test_momentum_*')
+
+# disable some creation ops
+backend_test.exclude('test_eyelike_*')
 
 # we only support float32
 backend_test.exclude('test_add_uint8_*')
 backend_test.exclude('test_sub_uint8_*')
 backend_test.exclude('test_div_uint8_*')
 backend_test.exclude('test_mul_uint8_*')
+backend_test.exclude('test_pow_types_int*')
 backend_test.exclude('test_cast_*')
 backend_test.exclude('test_castlike_*')
+backend_test.exclude('test_convinteger_*')
+backend_test.exclude('test_matmulinteger_*')
+
+# we don't support rounding
+backend_test.exclude('test_round_*')
+backend_test.exclude('test_ceil_*')
+
+# we don't support indexes
+backend_test.exclude('test_argmax_*')
+backend_test.exclude('test_argmin_*')
+backend_test.exclude('test_nonzero_*')
 
 # no support for nan or inf
 backend_test.exclude('test_isinf_*')
@@ -94,25 +113,30 @@ backend_test.exclude('test_scatternd_*')
 backend_test.exclude('test_dequantizelinear_*')
 backend_test.exclude('test_dynamicquantizelinear_*')
 backend_test.exclude('test_qlinearmatmul_*')
+backend_test.exclude('test_qlinearconv_*')
 backend_test.exclude('test_quantizelinear_*')
 
+# no rnn
+backend_test.exclude('test_gru_*')
+backend_test.exclude('test_rnn_*')
+backend_test.exclude('test_lstm_*')
+
+# no control flow
+backend_test.exclude('test_if_*')
+backend_test.exclude('test_loop*')
+
 # unsupported (strange) ops
-backend_test.exclude('test_argmax_*')
-backend_test.exclude('test_argmin_*')
 backend_test.exclude('test_bitwise_*')
 backend_test.exclude('test_blackmanwindow_*')
 backend_test.exclude('test_bernoulli_*')
 backend_test.exclude('test_cumsum_*')
 backend_test.exclude('test_tril_*')
 backend_test.exclude('test_triu_*')
-backend_test.exclude('test_convinteger_*')
 backend_test.exclude('test_col2im_*')
 backend_test.exclude('test_hammingwindow_*')
 backend_test.exclude('test_hannwindow_*')
 backend_test.exclude('test_hardmax_*')
-backend_test.exclude('test_gru_*')
 backend_test.exclude('test_gridsample_*')
-backend_test.exclude('test_if_*')
 backend_test.exclude('test_compress_*')
 backend_test.exclude('test_det_*')
 backend_test.exclude('test_dft_*')
@@ -124,16 +148,16 @@ backend_test.exclude('test_sequence_*')
 backend_test.exclude('test_nonmaxsuppression_*')
 backend_test.exclude('test_reversesequence_*')
 backend_test.exclude('test_roialign_*')
-backend_test.exclude('test_rnn_*')
 backend_test.exclude('test_top_k_*')
+backend_test.exclude('test_tfidfvectorizer_*')
+backend_test.exclude('test_stft_*')
+backend_test.exclude('test_melweightmatrix_*')
 
 # disable model tests for now since they are slow
 for x in backend_test.test_suite:
   if 'OnnxBackendRealModelTest' in str(type(x)):
     backend_test.exclude(str(x).split(" ")[0])
 
-#backend_test.include('test_tile_*')
-
 # passing node tests
 """
 backend_test.include('test_unsqueeze_*')
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index a9fa58001b..ef91dfefe0 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -133,7 +133,7 @@ class Tensor:
   def randn(cls, *shape, **kwargs): return cls(Tensor._rng.standard_normal(size=shape, dtype=np.float32), **kwargs)
 
   @classmethod
-  def arange(cls, stop, start=0, **kwargs): return cls(np.arange(start=start, stop=stop, dtype=np.float32), **kwargs)
+  def arange(cls, stop, start=0, step=1, **kwargs): return cls(np.arange(start=start, stop=stop, step=step, dtype=np.float32), **kwargs)
 
   # TODO: uniform should be a late binding thing
   # Return random number between -1 and 1