Tfjs: Support Rank-5 Tensors.

[nsthorat]

We've gotten a couple requests for this now.

What has to happen:

• Add a Tensor5D type: https://github.com/tensorflow/tfjs-core/blob/master/src/tensor.ts#L927
• Add a tf.tensor5d constructor: https://github.com/tensorflow/tfjs-core/blob/master/src/ops/array_ops.ts#L231
• Add 5D get output coords method, similar to 4D: https://github.com/tensorflow/tfjs-core/blob/master/src/kernels/webgl/shader_compiler.ts#L113
  https://github.com/tensorflow/tfjs-core/blob/master/src/kernels/webgl/shader_compiler.ts#L329
• Add 5D sampler similar to 4D sampler: https://github.com/tensorflow/tfjs-core/blob/master/src/kernels/webgl/shader_compiler.ts#L524
• Add coords dtype. It will likely have to be an array of vec's: https://github.com/tensorflow/tfjs-core/blob/master/src/kernels/webgl/shader_compiler.ts#L712

[zaidalyafeai]

@nsthorat, we also need it for https://github.com/tensorflow/tfjs/issues/250. I tried creating 5D tensors and an error was thrown
Error: GPU for rank 5 is not yet supported in src/kernels/webgl/shader_compiler.js (line 641).
I know how to solve it now thanks to your list. I can work on that if needed ...

[zaidalyafeai]

@nsthorat I need some pointers here. As of my understanding shaders language doesn't have vectors of size more than 4 ?

[nsthorat]

You can use a size-2 array of vec3s for the coords!

@zaidalyafeai are you going to take this up? I think @jgartman said he would also be interested in this since he needs this and 6D Tensors for another feature. Might be good to coordinate.

[zaidalyafeai]

@nsthorat, This is my first time working on a shader language. I can work on it but I'll need some help understanding the Sampler method. If @jgartman, wants to work on it, it is OK for me.

[jgartman]

@zaidalyafeai Like Nikhil said, I'm working on an op that might possibly require 6D tensors. I'm certainly no expert but I think if we read the existing code and WebGL documentation we can figure this out.

[zaidalyafeai]

@jgartman, I cant find any good resources on the internet. Could you share some on WebGL?

[jgartman]

@zaidalyafeai I thought this one was pretty good though there's a lot of info that isn't necessarily relevant to this task. So the first thing that was mentioned above was to define an output coords method. As far as I can tell, what this needs to do is take coordinates in the texture space and project them into the tensor space. The textures are what store the data for a tensor and I believe they are always 2D so to define a 5D output coords method you would have to project from 2D to 5D (I'm learning this as I go so hopefully all this info is accurate). I think it's a little clearer when you look at the 1D and 2D output coords methods so a good place to get started would be to try and understand those.

[zaidalyafeai]

@nsthorat Not sure if I can return an array of vectors from a function. Is that correct ?

[jgartman]

@zaidalyafeai Looks like that's not possible https://stackoverflow.com/questions/29734341/first-class-array-not-supported. One thing you might want to try is to stub out all the necessary methods and then try and create a 5D tensor and do an op with it. You'll get a bunch of errors but I think iteratively working through them will make development easier. Anyway, when I did that and declared the getOutputCoords function with an array return type WebGL threw the error 'ERROR: 0:91: 'first-class array' : not supported'. I think you can pass in an array and then modify its contents in the function so that might be something to look into.

[nsthorat]

Check out the response here, this is another way it can be done:

https://stackoverflow.com/questions/34030684/how-to-declare-a-function-that-returns-an-array-in-glsl-es-version-100

[zaidalyafeai]

Possible implementation of Sampler5D. Can anybody review it especially line 668

[nsthorat]

@zaidalyafeai if you can open a pull request I can take a closer look (hard to review an image).

Also, unit tests are the best way to test correctness!

[zaidalyafeai]

@nsthorat , sure will open one today when I have some free time. Do I have to create a test file for the shader_compiler.ts file ?

[nsthorat]

The unit test can call a unary op (relu or square or something) a rank-5 tensor.

[zaidalyafeai]

@nsthorat, It seems functions like get output coords are not used for unary operations. When I call the following
const a = tf.tensor5d([2], [1, 1, 1, 1, 1]); a.square();
I get no errors even though getOutput5DCoords() is not implemented correctly ( I made it return void). However, when I use conv2d on tensor4d it seems getOutput4DCoords() is called. Also I realized that the shader code of tensor4d expects getOutput4DCoords() to return a vector as I can see from the use of swizzling
ivec4 coords = getOutputCoords();
int batch = coords[0];
int d2 = coords[3];
ivec2 xRCCorner = coords.yz * strides - pads;
Guess that makes generalizing to tensor5d a little bit complicated ?

[jgartman]

@zaidalyafeai You may want to check your code again, I believe a unary op (as well as any op being executed on the GPU) will need that function. From kernels/webgl/unary_op.gpu:

void main() {
        float x = getAAtOutCoords();
        float y = unaryOperation(x);

        setOutput(y);
      }

[zaidalyafeai]

@jgartman , I made sure that I printed out the source code, getOutput4DCoords() is not called. getAAtOutCoords() calls texture2D.

[zaidalyafeai]

@nsthorat, I don't know how I missed that !! We can easily made a new data type called ivec5. We define that in the SHADER_PREFEX #L232 as a strut type.
struct ivec5 { int x; int y; int z; int w; int f; };
This allows creating arbitrary size tensors ...

[zaidalyafeai]

Created a PR https://github.com/tensorflow/tfjs-core/pull/1022

[nsthorat]

Thanks for all the hard work!