def METHOD_NAME(
    dim: int, device: str, placement, sbp
):
    shape = shapes[dim]
    y_shape = (shape[0],) if dim == 2 else (shape[0], *shape[2:])
    x_np = np.random.uniform(0, 1, shape)
    y_np = np.random.randint(0, shape[1], y_shape)
    def f(x, requires_grad):
        t = flow.tensor(x, device=device, requires_grad=requires_grad).to_global(
            placement=placement, sbp=[sbp]
        )
        if requires_grad:
            t.retain_grad()
        return t
    return f(x_np, True), f(y_np, False), f(x_np, True), f(y_np, False)