@thi.ng/tensors/defopn.js

0D/1D/2D/3D/4D tensors with extensible polymorphic operations and customizable storage

import { top } from "./top.js";
const defOpN = (fn) => {
  const f0 = (out, a) => {
    out.data[out.offset] = fn(a);
    return out;
  };
  const f1 = (out, a) => {
    const {
      data,
      offset,
      shape: [sx],
      stride: [tx]
    } = out;
    for (let x = 0; x < sx; x++) data[offset + x * tx] = fn(a);
    return out;
  };
  const f2 = (out, a) => {
    const {
      data,
      shape: [sx, sy],
      stride: [tx, ty],
      offset
    } = out;
    let ox;
    for (let x = 0; x < sx; x++) {
      ox = offset + x * tx;
      for (let y = 0; y < sy; y++) data[ox + y * ty] = fn(a);
    }
    return out;
  };
  const f3 = (out, a) => {
    const {
      data,
      shape: [sx, sy, sz],
      stride: [tx, ty, tz],
      offset
    } = out;
    let ox, oy;
    for (let x = 0; x < sx; x++) {
      ox = offset + x * tx;
      for (let y = 0; y < sy; y++) {
        oy = ox + y * ty;
        for (let z = 0; z < sz; z++) data[oy + z * tz] = fn(a);
      }
    }
    return out;
  };
  const f4 = (out, a) => {
    const {
      data,
      shape: [sx, sy, sz, sw],
      stride: [tx, ty, tz, tw],
      offset
    } = out;
    let ox, oy, oz;
    for (let x = 0; x < sx; x++) {
      ox = offset + x * tx;
      for (let y = 0; y < sy; y++) {
        oy = ox + y * ty;
        for (let z = 0; z < sz; z++) {
          oz = oy + z * tz;
          for (let w = 0; w < sw; w++) data[oz + w * tw] = fn(a);
        }
      }
    }
    return out;
  };
  return top(0, f0, f1, f2, f3, f4);
};
export {
  defOpN
};