code: def explicit_shard_W_axis_1(X, W):
      Xs = [X.to(d0), X.to(d1)]
      Ws = [
          W[:, : N // 2].to(d0),
          W[:, N // 2 :].to(d1),
      ]  # TODO: these shouldn't make copies on the original device
      # pad them to form the correct size
      Ws = [
          Ws[0].pad((None, (0, N // 2))),
          Ws[1].pad((None, (N // 2, 0))),
      ]
      for x in Xs:
          assert x.shape == X.shape
      for w in Ws:
          assert w.shape == W.shape
  
      # TODO: it shouldn't be faster with these realize
      for x in Xs + Ws:
          x.realize()
  
      def lm(x: Tensor, w: Tensor):
          # these are movement ops on the local device
          x = x.reshape(N, 1, N).expand(N, N, N)
          w = w.T.reshape(1, N, N).expand(N, N, N)
          m = x * w
          assert m.lazydata.st.views[0].mask is not None
          ret = m.sum(2)
          return ret
  
      # Os = [lm(Xs[0], Ws[0]), lm(Xs[1], Ws[1])]
      Os = [Xs[0] @ Ws[0], Xs[1] @ Ws[1]]
      for x in Os:
          x.realize()
      return Os[0].to(Device.DEFAULT) + Os[1].to(Device.DEFAULT)
  
imports: import numpy
  from tinygrad import Tensor, Device, GlobalCounters
  from tinygrad.helpers import Timing
label_desc: Use image node_id 1 for calling a global function with example usage: explicit_shard_W_axis_1(X, W) and returns: unknown, ret