| 39 | } |
| 40 | |
| 41 | getUserCode(): string { |
| 42 | let strideString; |
| 43 | if (this.sliceDim > 1) { |
| 44 | strideString = 'uniforms.strides[j]'; |
| 45 | } else { |
| 46 | strideString = 'uniforms.strides'; |
| 47 | } |
| 48 | const userCode = ` |
| 49 | ${main('index')} { |
| 50 | if (index < uniforms.size) { |
| 51 | let coords = getCoordsFromIndex(index); |
| 52 | var flattenIndex = 0; |
| 53 | for (var j = 0; j < uniforms.sliceDim; j = j + 1) { |
| 54 | let indexTemp = i32(round(getIndices(coords[0], j))); |
| 55 | let strideNum = ${strideString}; |
| 56 | flattenIndex = flattenIndex + indexTemp * strideNum; |
| 57 | } |
| 58 | |
| 59 | setOutputAtIndex(index, getA(flattenIndex, coords[1])); |
| 60 | } |
| 61 | } |
| 62 | `; |
| 63 | return userCode; |
| 64 | } |
| 65 | } |