3D Game of Life implementation in TensorFlow.js

gol.js

//mod function that handles negatives "properly"
function mod(a, b)
{
  return ((a%b)+b)%b
}

//shifts a tensor n times on the given axis
function shift(x, axis, n)
{
  n = mod(n, x.shape[axis]);
  x = x.split([n, x.shape[axis] - n], axis)
  x = tf.concat([x[1], x[0]], axis)
  return x;
}

//shifts a tensor the specified amounts on each axis
function multiShift(x, amounts)
{
  amounts.forEach((n, i) => x = shift(x, i, n));

  return x;
}

//3D game of life abcd implementation
//alive cells stay alive if a <= neighbors <= b
//dead cells come alive if c <= neighbors <= d
//you'll want to tidy this, eg:
//var newCells = tf.tidy(() => GOL(cells, 4, 5, 5, 5))
//cells.dispose()
//cells = newCells
function GOL(cells, a, b, c, d)
{
  //add some epsilon for sanity
  a = tf.scalar(a - 0.01);
  b = tf.scalar(b + 0.01);
  c = tf.scalar(c - 0.01);
  d = tf.scalar(d + 0.01);

  //count neighbors
  var neighbors = tf.zerosLike(cells);

  for(var x = -1; x <= 1; x++)
  for(var y = -1; y <= 1; y++)
  for(var z = -1; z <= 1; z++)
  {
    if(x == 0 && y == 0 && z == 0)
    {
      continue;
    }

    neighbors = neighbors.add(tf.where(
      multiShift(cells, [x, y, z]).greaterEqual(1),
      tf.onesLike(neighbors),
      tf.zerosLike(neighbors)
    ));

  }

  cells = tf.where(cells.greaterEqual(1),
    //if cell is currently living
    tf.where(tf.logicalAnd(
        neighbors.greaterEqual(a),
        neighbors.lessEqual(b)
      ),
      //if a <= neighbors <= b, stay alive
      tf.onesLike(cells),
      //else die
      tf.zerosLike(cells)
    ),
    //else cell is dead..
    tf.where(tf.logicalAnd(
        neighbors.greaterEqual(c),
        neighbors.lessEqual(d)
      ),
      //if c <= neighbors <=d, come alive
      tf.onesLike(cells),
      //else die
      tf.zerosLike(cells)
    )
  );

  return cells;
}