Adding benchmarks and testing for max op nanpropagation

benchmarks/python/single_ops.py

@@ -50,6 +50,11 @@ def time_maximum():
     mx.eval(a, b)
     time_fn(mx.maximum, a, b)
 
+def time_max():
+    a = mx.random.uniform(shape=(32, 1024, 1024))
+    a[1,1] = mx.nan
+    mx.eval(a)
+    time_fn(mx.max, a, 0)
 
 def time_negative():
     a = mx.random.uniform(shape=(10000, 1000))
@@ -108,6 +113,7 @@ if __name__ == "__main__":
 
     time_add()
     time_matmul()
+    time_max()
     time_maximum()
     time_exp()
     time_negative()