def METHOD_NAME(cls, ctx, op):
    axis = cls.get_axis(op.axis)
    a = ctx[op.inputs[0].key]
    if not isinstance(a, (list, tuple)):
        (inp,), device_id, xp = as_same_device(
            [a], device=op.device, ret_extra=True
        )
        with device(device_id):
            ctx[op.outputs[0].key] = xp.mean(
                inp, axis=axis, dtype=op.dtype, keepdims=bool(op.keepdims)
            )
    else:
        (_data, _count), device_id, xp = as_same_device(
            a, device=op.device, ret_extra=True
        )
        with device(device_id):
            chunk_count = xp.sum(
                _count, axis=axis, dtype=op.dtype, keepdims=bool(op.keepdims)
            )
            chunk_sum = xp.sum(
                _data, axis=axis, dtype=op.dtype, keepdims=bool(op.keepdims)
            )
            ctx[op.outputs[0].key] = xp.true_divide(
                chunk_sum, chunk_count, dtype=op.dtype
            )