From 792a494eb89e4b12702180a558121a3bbccfba4a Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Wed, 22 May 2024 20:43:21 -0400
Subject: [PATCH] fix various examples (#4691)

* fix examples that used ax1 and ax2 for transpose

* fix that

* update those
---
 examples/compile_tensorflow.py      | 1 +
 examples/simple_conv_bn.py          | 2 --
 examples/so_vits_svc.py             | 2 +-
 examples/vits.py                    | 4 ++--
 {examples => extra}/f16_w_uint32.py | 0
 5 files changed, 4 insertions(+), 5 deletions(-)
 rename {examples => extra}/f16_w_uint32.py (100%)

diff --git a/examples/compile_tensorflow.py b/examples/compile_tensorflow.py
index 4fee2a41c2..e3def3bbdb 100644
--- a/examples/compile_tensorflow.py
+++ b/examples/compile_tensorflow.py
@@ -2,6 +2,7 @@
 
 import os, sys
 os.environ["CLANG"] = '1'
+os.environ["JIT"] = '2'
 
 import numpy as np
 import subprocess
diff --git a/examples/simple_conv_bn.py b/examples/simple_conv_bn.py
index 7d5add4da7..287691016f 100644
--- a/examples/simple_conv_bn.py
+++ b/examples/simple_conv_bn.py
@@ -1,5 +1,3 @@
-# to start thinking about the $2,000 norm fusion bounty
-
 from tinygrad.tensor import Tensor
 from tinygrad.nn import Conv2d, BatchNorm2d
 from tinygrad.nn.state import get_parameters
diff --git a/examples/so_vits_svc.py b/examples/so_vits_svc.py
index 4829820c9a..b8872a04ce 100644
--- a/examples/so_vits_svc.py
+++ b/examples/so_vits_svc.py
@@ -223,7 +223,7 @@ class ConvFeatureExtractionModel:
         return conv
       assert (is_layer_norm and is_group_norm) == False, "layer norm and group norm are exclusive"
       if is_layer_norm:
-        return [make_conv(), partial(Tensor.dropout, p=dropout),[partial(Tensor.transpose, ax1=-2, ax2=-1), nn.LayerNorm(dim, elementwise_affine=True), partial(Tensor.transpose, ax1=-2, ax2=-1)], Tensor.gelu]
+        return [make_conv(), partial(Tensor.dropout, p=dropout),[partial(Tensor.transpose, dim0=-2, dim1=-1), nn.LayerNorm(dim, elementwise_affine=True), partial(Tensor.transpose, dim0=-2, dim1=-1)], Tensor.gelu]
       elif is_group_norm and mode == "default":
         return [make_conv(), partial(Tensor.dropout, p=dropout), nn.GroupNorm(dim, dim, affine=True), Tensor.gelu]
       elif is_group_norm and mode == "group_norm_masked":
diff --git a/examples/vits.py b/examples/vits.py
index 5f04a41e6b..d96152dd0d 100644
--- a/examples/vits.py
+++ b/examples/vits.py
@@ -489,14 +489,14 @@ def split(tensor, split_sizes, dim=0):  # if split_sizes is an integer, convert
     start += size
   return [tensor.slice(s) for s in slices]
 def gather(x, indices, axis):
-  indices = (indices < 0).where(indices + x.shape[axis], indices).transpose(ax1=axis, ax2=0)
+  indices = (indices < 0).where(indices + x.shape[axis], indices).transpose(0, axis)
   permute_args = list(range(x.ndim))
   permute_args[0], permute_args[axis] = permute_args[axis], permute_args[0]
   permute_args.append(permute_args.pop(0))
   x = x.permute(*permute_args)
   reshape_arg = [1] * x.ndim + [x.shape[-1]]
   return ((indices.unsqueeze(indices.ndim).expand(*indices.shape, x.shape[-1]) ==
-           Tensor.arange(x.shape[-1]).reshape(*reshape_arg).expand(*indices.shape, x.shape[-1])) * x).sum(indices.ndim).transpose(ax1=0, ax2=axis)
+           Tensor.arange(x.shape[-1]).reshape(*reshape_arg).expand(*indices.shape, x.shape[-1])) * x).sum(indices.ndim).transpose(0, axis)
 
 def norm_except_dim(v, dim):
   if dim == -1: return np.linalg.norm(v)
diff --git a/examples/f16_w_uint32.py b/extra/f16_w_uint32.py
similarity index 100%
rename from examples/f16_w_uint32.py
rename to extra/f16_w_uint32.py