```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Fused TripletMarginLoss CUDA kernel (p=2, swap=False, reduction='mean'):
#   d_pos(i) = sqrt( sum_j (a_ij - p_ij + eps)^2 )
#   d_neg(i) = sqrt( sum_j (a_ij - n_ij + eps)^2 )
#   loss     = mean_i( max(d_pos(i) - d_neg(i) + margin, 0) )
# One thread block per sample computes both squared distances with
# vectorized float4 loads, reduces them in shared memory / shuffles,
# applies the hinge, and accumulates the scaled loss in double precision.
_triplet_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

__device__ __forceinline__ void block_emit_loss(float sp, float sn,
                                                double margin, double scale,
                                                double* out) {
    __shared__ double s_pos[32];
    __shared__ double s_neg[32];
    double dsp = (double)sp;
    double dsn = (double)sn;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        dsp += __shfl_down_sync(0xffffffffu, dsp, o);
        dsn += __shfl_down_sync(0xffffffffu, dsn, o);
    }
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) {
        s_pos[warp] = dsp;
        s_neg[warp] = dsn;
    }
    __syncthreads();
    if (warp == 0) {
        const int nw = (blockDim.x + 31) >> 5;
        dsp = (lane < nw) ? s_pos[lane] : 0.0;
        dsn = (lane < nw) ? s_neg[lane] : 0.0;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            dsp += __shfl_down_sync(0xffffffffu, dsp, o);
            dsn += __shfl_down_sync(0xffffffffu, dsn, o);
        }
        if (lane == 0) {
            const double dist_pos = sqrt(dsp);
            const double dist_neg = sqrt(dsn);
            double loss = dist_pos - dist_neg + margin;
            if (loss < 0.0) loss = 0.0;
            atomicAdd(out, loss * scale);
        }
    }
}

__global__ void triplet_margin_loss_vec4_kernel(
    const float4* __restrict__ anchor,
    const float4* __restrict__ positive,
    const float4* __restrict__ negative,
    double* __restrict__ out,
    const double margin,
    const float eps,
    const double scale,
    const int vec_size)
{
    const int64_t row = blockIdx.x;
    const float4* a_row = anchor + row * (int64_t)vec_size;
    const float4* p_row = positive + row * (int64_t)vec_size;
    const float4* n_row = negative + row * (int64_t)vec_size;

    float sp = 0.0f;
    float sn = 0.0f;
    for (int i = threadIdx.x; i < vec_size; i += blockDim.x) {
        const float4 a = __ldg(a_row + i);
        const float4 p = __ldg(p_row + i);
        const float4 n = __ldg(n_row + i);
        const float dp0 = a.x - p.x + eps;
        const float dp1 = a.y - p.y + eps;
        const float dp2 = a.z - p.z + eps;
        const float dp3 = a.w - p.w + eps;
        const float dn0 = a.x - n.x + eps;
        const float dn1 = a.y - n.y + eps;
        const float dn2 = a.z - n.z + eps;
        const float dn3 = a.w - n.w + eps;
        sp += dp0 * dp0 + dp1 * dp1 + dp2 * dp2 + dp3 * dp3;
        sn += dn0 * dn0 + dn1 * dn1 + dn2 * dn2 + dn3 * dn3;
    }
    block_emit_loss(sp, sn, margin, scale, out);
}

__global__ void triplet_margin_loss_scalar_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    double* __restrict__ out,
    const double margin,
    const float eps,
    const double scale,
    const int64_t features)
{
    const int64_t row = blockIdx.x;
    const float* a_row = anchor + row * features;
    const float* p_row = positive + row * features;
    const float* n_row = negative + row * features;

    float sp = 0.0f;
    float sn = 0.0f;
    for (int64_t i = threadIdx.x; i < features; i += blockDim.x) {
        const float a = __ldg(a_row + i);
        const float p = __ldg(p_row + i);
        const float n = __ldg(n_row + i);
        const float dp = a - p + eps;
        const float dn = a - n + eps;
        sp += dp * dp;
        sn += dn * dn;
    }
    block_emit_loss(sp, sn, margin, scale, out);
}

torch::Tensor triplet_margin_loss_cuda(torch::Tensor anchor, torch::Tensor positive,
                                       torch::Tensor negative, double margin, double eps) {
    TORCH_CHECK(anchor.is_cuda() && positive.is_cuda() && negative.is_cuda(),
                "inputs must be CUDA tensors");
    TORCH_CHECK(anchor.scalar_type() == torch::kFloat &&
                positive.scalar_type() == torch::kFloat &&
                negative.scalar_type() == torch::kFloat,
                "inputs must be float32");

    auto a = anchor.contiguous();
    auto p = positive.contiguous();
    auto n = negative.contiguous();

    const int64_t batch = a.size(0);
    const int64_t features = batch > 0 ? a.numel() / batch : 0;

    auto out = torch::zeros({}, torch::TensorOptions().dtype(torch::kDouble).device(a.device()));

    if (batch > 0 && features > 0) {
        const int threads = 256;
        const int blocks = (int)batch;
        const double scale = 1.0 / (double)batch;
        cudaStream_t stream = at::cuda::getCurrentCUDAStream();
        if ((features & 3) == 0) {
            triplet_margin_loss_vec4_kernel<<<blocks, threads, 0, stream>>>(
                reinterpret_cast<const float4*>(a.data_ptr<float>()),
                reinterpret_cast<const float4*>(p.data_ptr<float>()),
                reinterpret_cast<const float4*>(n.data_ptr<float>()),
                out.data_ptr<double>(),
                margin, (float)eps, scale, (int)(features >> 2));
        } else {
            triplet_margin_loss_scalar_kernel<<<blocks, threads, 0, stream>>>(
                a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
                out.data_ptr<double>(), margin, (float)eps, scale, features);
        }
    }
    return out.to(torch::kFloat);
}
"""

_triplet_cpp_src = (
    "torch::Tensor triplet_margin_loss_cuda(torch::Tensor anchor, torch::Tensor positive, "
    "torch::Tensor negative, double margin, double eps);"
)

_triplet_ext = load_inline(
    name="triplet_margin_loss_fused",
    cpp_sources=_triplet_cpp_src,
    cuda_sources=_triplet_cuda_src,
    functions=["triplet_margin_loss_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Triplet Margin Loss accelerated with a single fused CUDA kernel that
    computes both pairwise distances, the hinge, and the mean reduction
    in one pass over the inputs.
    """
    def __init__(self, margin=1.0):
        super(ModelNew, self).__init__()
        self.margin = float(margin)
        self.eps = 1e-6  # matches torch.nn.TripletMarginLoss default

    def forward(self, anchor, positive, negative):
        if anchor.is_cuda and anchor.dtype == torch.float32:
            return _triplet_ext.triplet_margin_loss_cuda(
                anchor, positive, negative, self.margin, self.eps)
        return F.triplet_margin_loss(anchor, positive, negative,
                                     margin=self.margin, p=2.0, eps=self.eps,
                                     swap=False, reduction="mean")
```