@Jeremy if you think this or part of it belongs in harebrain, feel free to move it over.
Big thanks to @sgugger for finding the Path and Just Swift libraries. I definitely looked over your shoulder on github to get a little help figuring those out.
A few things that did not go as I expected:
Broadcasting in s4tf with Tensors does not play super nice with operators like >, <, >=, <=, ==, !=. var a = Tensor([10.0, 6.0, -4.0]) print(a > 0) returns false
Note: np/pytorch return [1, 1, 0] Edit: Thanks @jekbradbury for pointing out that S4TF does have this ability. I simply needed to use the .>, .<=, etc. operators prefixed with the period.
Either I could not find it or S4TF does not have a sum by axis (nor a min/max with an axis parameter) like np.sum(a, axis=0). If I overlooked it, please share some knowledge!
To get around this I converted the TensorFlow Tensors to np and back. Edit: Thanks @sgugger and @dan-zheng for pointing out that there are squeezingAxes and alongAxes arguments to the S4TF Tensor.sum method that I overlooked. Looks like it is documented here under the “Extensions” heading.
If you scroll down to the bottom you can see that S4TF’s matrix multiply is faster than the PyTorch one in Jeremy’s notebook. However that is not really a fair comparison (only comparing a single operation, not using the same hardware, etc).
I’ll try to do a more interesting comparison when I do the next notebook.
Thanks for putting this together and sharing it, Jeff.
Minor issue: in your first %%timeit cell with slowMatMul, the mean is greater than the max. After diving into the code, I think the issue is that you run the loop n_times+1 in the timeitMagic function. See my PR below.