```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
from torch.utils.cpp_extension import load_inline

# Fused TripletMarginLoss CUDA kernel.
# Computes, for each row i:
#   dist_pos = ||anchor_i - positive_i + eps||_2
#   dist_neg = ||anchor_i - negative_i + eps||_2
#   loss_i   = max(dist_pos - dist_neg + margin, 0)
# and returns the mean loss (reduction='mean', p=2, swap=False).
triplet_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__inline__ __device__ double warp_reduce_sum(double v) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffff, v, offset);
    }
    return v;
}

__inline__ __device__ double block_reduce_sum(double v, double* sdata) {
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    v = warp_reduce_sum(v);
    if (lane == 0) sdata[wid] = v;
    __syncthreads();
    int nwarp = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        v = (lane < nwarp) ? sdata[lane] : 0.0;
        v = warp_reduce_sum(v);
        if (lane == 0) sdata[0] = v;
    }
    __syncthreads();
    double out = sdata[0];
    __syncthreads();
    return out;
}

// Vectorized path (row_size % 4 == 0): one block per row.
__global__ void triplet_margin_loss_vec4_kernel(
        const float* __restrict__ anchor,
        const float* __restrict__ positive,
        const float* __restrict__ negative,
        double* __restrict__ losses,
        long row_size,
        float margin,
        float eps) {
    long row = blockIdx.x;
    const float4* a4 = reinterpret_cast<const float4*>(anchor   + row * row_size);
    const float4* p4 = reinterpret_cast<const float4*>(positive + row * row_size);
    const float4* n4 = reinterpret_cast<const float4*>(negative + row * row_size);
    long nvec = row_size >> 2;

    double sp = 0.0;
    double sn = 0.0;
    for (long i = threadIdx.x; i < nvec; i += blockDim.x) {
        float4 av = a4[i];
        float4 pv = p4[i];
        float4 nv = n4[i];
        float d;
        d = av.x - pv.x + eps; sp += (double)d * (double)d;
        d = av.y - pv.y + eps; sp += (double)d * (double)d;
        d = av.z - pv.z + eps; sp += (double)d * (double)d;
        d = av.w - pv.w + eps; sp += (double)d * (double)d;
        d = av.x - nv.x + eps; sn += (double)d * (double)d;
        d = av.y - nv.y + eps; sn += (double)d * (double)d;
        d = av.z - nv.z + eps; sn += (double)d * (double)d;
        d = av.w - nv.w + eps; sn += (double)d * (double)d;
    }

    __shared__ double sdata_p[32];
    __shared__ double sdata_n[32];
    sp = block_reduce_sum(sp, sdata_p);
    sn = block_reduce_sum(sn, sdata_n);

    if (threadIdx.x == 0) {
        double l = sqrt(sp) - sqrt(sn) + (double)margin;
        losses[row] = l > 0.0 ? l : 0.0;
    }
}

// Generic fallback path (any row_size): one block per row.
__global__ void triplet_margin_loss_scalar_kernel(
        const float* __restrict__ anchor,
        const float* __restrict__ positive,
        const float* __restrict__ negative,
        double* __restrict__ losses,
        long row_size,
        float margin,
        float eps) {
    long row = blockIdx.x;
    const float* a = anchor   + row * row_size;
    const float* p = positive + row * row_size;
    const float* n = negative + row * row_size;

    double sp = 0.0;
    double sn = 0.0;
    for (long i = threadIdx.x; i < row_size; i += blockDim.x) {
        float av = a[i];
        float dp = av - p[i] + eps;
        float dn = av - n[i] + eps;
        sp += (double)dp * (double)dp;
        sn += (double)dn * (double)dn;
    }

    __shared__ double sdata_p[32];
    __shared__ double sdata_n[32];
    sp = block_reduce_sum(sp, sdata_p);
    sn = block_reduce_sum(sn, sdata_n);

    if (threadIdx.x == 0) {
        double l = sqrt(sp) - sqrt(sn) + (double)margin;
        losses[row] = l > 0.0 ? l : 0.0;
    }
}

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(),
                "All inputs must be CUDA tensors.");
    auto a = anchor.contiguous();
    auto p = positive.contiguous();
    auto n = negative.contiguous();
    TORCH_CHECK(a.scalar_type() == torch::kFloat32 &&
                p.scalar_type() == torch::kFloat32 &&
                n.scalar_type() == torch::kFloat32,
                "Inputs must be float32.");
    TORCH_CHECK(a.numel() == p.numel() && a.numel() == n.numel(),
                "Input shapes must match.");

    long B = a.size(0);
    long D = (a.dim() > 1) ? (a.numel() / B) : 1;

    auto losses = torch::empty({B}, a.options().dtype(torch::kFloat64));

    const int block = 256;
    dim3 grid((unsigned int)B);
    auto stream = at::cuda::getCurrentCUDAStream();

    if ((D & 3) == 0) {
        triplet_margin_loss_vec4_kernel<<<grid, block, 0, stream>>>(
            a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
            losses.data_ptr<double>(), D, (float)margin, (float)eps);
    } else {
        triplet_margin_loss_scalar_kernel<<<grid, block, 0, stream>>>(
            a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
            losses.data_ptr<double>(), D, (float)margin, (float)eps);
    }

    // Deterministic double-precision mean, then cast back to float32 to
    // match the dtype of torch.nn.TripletMarginLoss(reduction='mean').
    return losses.mean().to(torch::kFloat32);
}
"""

triplet_cpp_source = (
    "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_source,
    cuda_sources=triplet_cuda_source,
    functions=["triplet_margin_loss_cuda"],
    verbose=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized Triplet Margin Loss (margin parameter, p=2, eps=1e-6,
    swap=False, reduction='mean') implemented as a single fused CUDA kernel
    plus a tiny final mean reduction.
    """
    def __init__(self, margin=1.0):
        super(ModelNew, self).__init__()
        self.margin = float(margin)
        self.eps = 1e-6  # default of torch.nn.TripletMarginLoss
        self.ext = _triplet_ext

    def forward(self, anchor, positive, negative):
        return self.ext.triplet_margin_loss_cuda(
            anchor, positive, negative, self.margin, self.eps
        )


batch_size = 32768
input_shape = (8192,)
dim = 1


def get_inputs():
    scale = torch.rand(())
    return [
        torch.rand(batch_size, *input_shape).cuda() * scale.cuda(),
        torch.rand(batch_size, *input_shape).cuda(),
        torch.rand(batch_size, *input_shape).cuda(),
    ]


def get_init_inputs():
    return [1.0]  # Default margin
```