from tinygrad import Tensor, dtypes
import numpy as np
def length_masked_ce_loss(model, inputs, targets, lengths):
  # Run model on inputs
  logits = model(inputs).cast(dtypes.float32).contiguous()

  # Mask padding tokens
  length_mask = Tensor(np.arange(inputs.shape[1])[None, :] < lengths[:, None], requires_grad=False)

  # Calculate the loss
  ce = logits.sparse_categorical_crossentropy(Tensor(targets, requires_grad=False)).mul(length_mask)
  loss = ce.sum() / length_mask.sum()
  return loss