Help for MLX equivalent of PyTorch Code

[sachinraja13]

I need help to port this piece of code to MLX ensuring that gradients are correctly computed using nn.value_and_grad.

output_class = class_embed(memory) # Linear layer
ref_token_index = torch.topk(output_class(-1)[0], topk, dim=1)[1] # bs, nq
ref_token_coord = torch.gather(output_proposals, 1, ref_token_index.unsqueeze(-1).repeat(1, 1, 4))

class_embed is a linear layer which needs gradients to be updated. output_proposals and memory are computed using parameters that need gradients to be computed for updation.

This is what I have done:

output_class = class_embed(memory)
ref_token_index = mx.argpartition(output_class(axis=-1) * -1, topk , axis=1)[:, :topk] 
ref_token_coord = output_proposals[mx.arange(batch_size)[:, None, None], ref_token_index[..., None], mx.arange(4)[None, None, :]]

I'm facing this error during back prop:
ValueError: [Primitive::vjp] Not implemented for ArgPartition.