Shape.Broadcasting.cs

using System;
using NumSharp.Utilities;

namespace NumSharp
{
    /// <summary>
    ///     Broadcasting operations for Shape.
    /// </summary>
    /// <remarks>
    ///     Based on https://numpy.org/doc/stable-1.16.1/user/basics.broadcasting.html
    ///     These are pure shape operations - they compute output shapes without touching array data.
    /// </remarks>
    public partial struct Shape
    {
        /// <summary>
        ///     Resolves the output shape when broadcasting two shapes together.
        /// </summary>
        /// <param name="leftShape">First shape to broadcast.</param>
        /// <param name="rightShape">Second shape to broadcast.</param>
        /// <returns>The resulting broadcast shape.</returns>
        /// <exception cref="IncorrectShapeException">If shapes cannot be broadcast together.</exception>
        public static Shape ResolveReturnShape(Shape leftShape, Shape rightShape)
        {
            if (leftShape == rightShape)
                return leftShape;

            if (leftShape.IsScalar || leftShape.NDim == 1 && leftShape[0] == 1)
                return rightShape;

            if (rightShape.IsScalar || rightShape.NDim == 1 && rightShape[0] == 1)
                return leftShape;

            int i, nd, k;
            long tmp;

            nd = Math.Max(leftShape.NDim, rightShape.NDim);

            // This is the shared shape aka the target broadcast
            var mit = new long[nd];

            // Discover the broadcast shape in each dimension
            for (i = 0; i < nd; i++)
            {
                mit[i] = 1;

                // This prepends 1 to shapes not already equal to nd
                k = i + leftShape.NDim - nd;
                if (k >= 0)
                {
                    tmp = leftShape.dimensions[k];
                    if (tmp == 1)
                    {
                        goto _continue;
                    }

                    if (mit[i] == 1)
                    {
                        mit[i] = tmp;
                    }
                    else if (mit[i] != tmp)
                    {
                        throw new IncorrectShapeException("shape mismatch: objects cannot be broadcast to a single shape");
                    }
                }

                _continue:
                // This prepends 1 to shapes not already equal to nd
                k = i + rightShape.NDim - nd;
                if (k >= 0)
                {
                    tmp = rightShape.dimensions[k];
                    if (tmp == 1)
                    {
                        continue;
                    }

                    if (mit[i] == 1)
                    {
                        mit[i] = tmp;
                    }
                    else if (mit[i] != tmp)
                    {
                        throw new IncorrectShapeException("shape mismatch: objects cannot be broadcast to a single shape");
                    }
                }
            }

            // We don't need to mark it as IsBroadcasted because technically, it has not broadcasted strides.
            return mit; // implicit cast
        }

        /// <summary>
        ///     Resolves the output shape when broadcasting multiple shapes together.
        /// </summary>
        /// <param name="shapes">Shapes to broadcast.</param>
        /// <returns>The resulting broadcast shape.</returns>
        /// <exception cref="IncorrectShapeException">If shapes cannot be broadcast together.</exception>
        public static Shape ResolveReturnShape(params Shape[] shapes)
        {
            if (shapes.Length == 0)
                return default;

            if (shapes.Length == 1)
                return shapes[0].Clean();

            Shape mit;
            int i, nd, k;
            long tmp;

            int len = shapes.Length;

            // Discover the broadcast number of dimensions
            // Gets the largest ndim of all iterators
            nd = 0;
            for (i = 0; i < len; i++)
                nd = Math.Max(nd, shapes[i].NDim);

            // This is the shared shape aka the target broadcast
            mit = Shape.Empty(nd);

            // Discover the broadcast shape in each dimension
            for (i = 0; i < nd; i++)
            {
                mit.dimensions[i] = 1;
                for (int targetIndex = 0; targetIndex < len; targetIndex++)
                {
                    Shape target = shapes[targetIndex];
                    // This prepends 1 to shapes not already equal to nd
                    k = i + target.NDim - nd;
                    if (k >= 0)
                    {
                        tmp = target.dimensions[k];
                        if (tmp == 1)
                        {
                            continue;
                        }

                        if (mit.dimensions[i] == 1)
                        {
                            mit.dimensions[i] = tmp;
                        }
                        else if (mit.dimensions[i] != tmp)
                        {
                            throw new IncorrectShapeException("shape mismatch: objects cannot be broadcast to a single shape");
                        }
                    }
                }
            }

            // We don't need to mark it as IsBroadcasted because technically, it has not broadcasted strides.
            return mit.Clean();
        }

        /// <summary>
        ///     Broadcasts multiple shapes and returns the broadcasted shapes with computed strides.
        /// </summary>
        /// <param name="shapes">Shapes to broadcast.</param>
        /// <returns>Array of broadcasted shapes with stride=0 where dimensions were stretched.</returns>
        /// <exception cref="IncorrectShapeException">If shapes cannot be broadcast together.</exception>
        public static Shape[] Broadcast(params Shape[] shapes)
        {
            if (shapes.Length == 0)
                return Array.Empty<Shape>();

            if (shapes.Length == 1)
                return new Shape[] { shapes[0].Clean() };

            int i, nd, k, j;
            long tmp;

            var ret = new Shape[shapes.Length];
            int len = shapes.Length;

            tmp = 0;
            // Discover the broadcast number of dimensions
            nd = 0;
            for (i = 0; i < len; i++)
                nd = Math.Max(nd, shapes[i].NDim);

            // Use temporary array for broadcast dimensions (not Shape.Empty)
            var mitDims = new long[nd];

            // Discover the broadcast shape in each dimension
            for (i = 0; i < nd; i++)
            {
                mitDims[i] = 1;
                for (int targetIndex = 0; targetIndex < len; targetIndex++)
                {
                    Shape target = shapes[targetIndex];
                    k = i + target.NDim - nd;
                    if (k >= 0)
                    {
                        tmp = target.dimensions[k];
                        if (tmp == 1)
                            continue;

                        if (mitDims[i] == 1)
                            mitDims[i] = tmp;
                        else if (mitDims[i] != tmp)
                            throw new IncorrectShapeException("shape mismatch: objects cannot be broadcast to a single shape");
                    }
                }
            }

            // Create broadcast shapes for each input
            for (i = 0; i < len; i++)
            {
                Shape ogiter = shapes[i];
                int ogNd = ogiter.NDim;

                // Compute broadcast strides
                var broadcastStrides = new long[nd];
                for (j = 0; j < nd; j++)
                {
                    k = j + ogNd - nd;
                    if ((k < 0) || ogiter.dimensions[k] != mitDims[j])
                        broadcastStrides[j] = 0;
                    else
                        broadcastStrides[j] = ogiter.strides[k];
                }

                // Create immutable shape via constructor
                long bufSize = ogiter.bufferSize > 0 ? ogiter.bufferSize : ogiter.size;
                ret[i] = new Shape(
                    (long[])mitDims.Clone(),
                    broadcastStrides,
                    ogiter.offset,
                    bufSize
                );
            }

            return ret;
        }

        /// <summary>
        ///     Broadcasts two shapes and returns the broadcasted shapes with computed strides.
        /// </summary>
        /// <param name="leftShape">First shape to broadcast.</param>
        /// <param name="rightShape">Second shape to broadcast.</param>
        /// <returns>Tuple of (LeftBroadcastedShape, RightBroadcastedShape).</returns>
        /// <exception cref="IncorrectShapeException">If shapes cannot be broadcast together.</exception>
        public static (Shape LeftShape, Shape RightShape) Broadcast(Shape leftShape, Shape rightShape)
        {
            // Fast path: IDENTICAL shapes broadcast to themselves. The shape hash is only a cheap
            // pre-filter — it COLLIDES across dims that differ but hash alike (notably a 0-length
            // axis vs a size-1 axis: (1,1) and (0,1) both hash to the same value). Confirm the
            // dimensions are actually equal before short-circuiting, else e.g. broadcast_to((1,1),
            // (0,1)) wrongly returned (1,1) instead of stretching the size-1 axis to the 0 axis.
            if (leftShape._hashCode != 0 && leftShape._hashCode == rightShape._hashCode
                && leftShape.NDim == rightShape.NDim)
            {
                bool sameDims = true;
                for (int d = 0; d < leftShape.NDim; d++)
                    if (leftShape.dimensions[d] != rightShape.dimensions[d]) { sameDims = false; break; }
                if (sameDims)
                    return (leftShape, rightShape);
            }

            int i, nd, k, j;
            long tmp;

            // Is left a scalar / size-1 — broadcast to right's shape with zero strides. Guard on
            // rightShape.NDim >= leftShape.NDim: a 1-D [1] broadcast against a 0-D scalar must keep
            // its rank (result [1], not []), so the size-1 collapse only applies when the other
            // operand has at least as many dimensions. NumPy: result ndim == max(ndims).
            if ((leftShape.IsScalar || (leftShape.NDim == 1 && leftShape.size == 1)) && rightShape.NDim >= leftShape.NDim)
            {
                var zeroStrides = new long[rightShape.NDim];
                long leftBufSize = leftShape.bufferSize > 0 ? leftShape.bufferSize : leftShape.size;
                var left = new Shape(
                    (long[])rightShape.dimensions.Clone(),
                    zeroStrides,
                    leftShape.offset,
                    leftBufSize
                );
                return (left, rightShape);
            }

            // Is right a scalar / size-1 — broadcast to left's shape with zero strides. Symmetric
            // guard: leftShape.NDim >= rightShape.NDim so a 0-D right against a 1-D [1] left keeps
            // rank (result [1]).
            if ((rightShape.IsScalar || (rightShape.NDim == 1 && rightShape.size == 1)) && leftShape.NDim >= rightShape.NDim)
            {
                var zeroStrides = new long[leftShape.NDim];
                long rightBufSize = rightShape.bufferSize > 0 ? rightShape.bufferSize : rightShape.size;
                var right = new Shape(
                    (long[])leftShape.dimensions.Clone(),
                    zeroStrides,
                    rightShape.offset,
                    rightBufSize
                );
                return (leftShape, right);
            }

            // General case: compute broadcast shape
            nd = Math.Max(rightShape.NDim, leftShape.NDim);

            // Compute broadcast dimensions into temporary array
            var mitDims = new long[nd];
            for (i = 0; i < nd; i++)
            {
                mitDims[i] = 1;

                k = i + leftShape.NDim - nd;
                if (k >= 0)
                {
                    tmp = leftShape.dimensions[k];
                    if (tmp != 1)
                    {
                        if (mitDims[i] == 1)
                            mitDims[i] = tmp;
                        else if (mitDims[i] != tmp)
                            throw new IncorrectShapeException("shape mismatch: objects cannot be broadcast to a single shape");
                    }
                }

                k = i + rightShape.NDim - nd;
                if (k >= 0)
                {
                    tmp = rightShape.dimensions[k];
                    if (tmp != 1)
                    {
                        if (mitDims[i] == 1)
                            mitDims[i] = tmp;
                        else if (mitDims[i] != tmp)
                            throw new IncorrectShapeException("shape mismatch: objects cannot be broadcast to a single shape");
                    }
                }
            }

            // Compute left broadcast strides
            var leftStrides = new long[nd];
            for (j = 0; j < nd; j++)
            {
                k = j + leftShape.NDim - nd;
                if ((k < 0) || leftShape.dimensions[k] != mitDims[j])
                    leftStrides[j] = 0;
                else
                    leftStrides[j] = leftShape.strides[k];
            }

            // Compute right broadcast strides
            var rightStrides = new long[nd];
            for (j = 0; j < nd; j++)
            {
                k = j + rightShape.NDim - nd;
                if ((k < 0) || rightShape.dimensions[k] != mitDims[j])
                    rightStrides[j] = 0;
                else
                    rightStrides[j] = rightShape.strides[k];
            }

            // Create immutable shapes via constructors
            long leftBufSize2 = leftShape.bufferSize > 0 ? leftShape.bufferSize : leftShape.size;
            long rightBufSize2 = rightShape.bufferSize > 0 ? rightShape.bufferSize : rightShape.size;

            var leftResult = new Shape(mitDims, leftStrides, leftShape.offset, leftBufSize2);
            var rightResult = new Shape((long[])mitDims.Clone(), rightStrides, rightShape.offset, rightBufSize2);

            return (leftResult, rightResult);
        }

        /// <summary>
        ///     Checks if the given shapes can be broadcast together.
        /// </summary>
        /// <param name="shapes">Shapes to check.</param>
        /// <returns>True if all shapes are broadcastable, false otherwise.</returns>
        public static unsafe bool AreBroadcastable(params Shape[] shapes)
        {
            if (shapes.Length <= 1)
                return true;

            int i, nd, k;
            long tmp;

            int len = shapes.Length;

            // Discover the broadcast number of dimensions
            // Gets the largest ndim of all iterators
            nd = 0;
            for (i = 0; i < len; i++)
                nd = Math.Max(nd, shapes[i].NDim);

            // This is the shared shape aka the target broadcast
            var mit = stackalloc long[nd];

            // Discover the broadcast shape in each dimension
            for (i = 0; i < nd; i++)
            {
                mit[i] = 1;
                for (int targetIndex = 0; targetIndex < len; targetIndex++)
                {
                    Shape target = shapes[targetIndex];
                    // This prepends 1 to shapes not already equal to nd
                    k = i + target.NDim - nd;
                    if (k >= 0)
                    {
                        tmp = target.dimensions[k];
                        if (tmp == 1)
                        {
                            continue;
                        }

                        if (mit[i] == 1)
                        {
                            mit[i] = tmp;
                        }
                        else if (mit[i] != tmp)
                        {
                            return false;
                        }
                    }
                }
            }

            return true;
        }

        /// <summary>
        ///     Checks if the given dimension arrays can be broadcast together.
        /// </summary>
        /// <param name="shapes">Dimension arrays to check.</param>
        /// <returns>True if all shapes are broadcastable, false otherwise.</returns>
        public static unsafe bool AreBroadcastable(params int[][] shapes)
        {
            if (shapes.Length <= 1)
                return true;

            int i, nd, k, tmp;

            int len = shapes.Length;

            tmp = 0;
            // Discover the broadcast number of dimensions
            // Gets the largest ndim of all iterators
            nd = 0;
            for (i = 0; i < len; i++)
                nd = Math.Max(nd, shapes[i].Length);

            // This is the shared shape aka the target broadcast
            var mit = stackalloc int[nd];

            // Discover the broadcast shape in each dimension
            for (i = 0; i < nd; i++)
            {
                mit[i] = 1;
                for (int targetIndex = 0; targetIndex < len; targetIndex++)
                {
                    int[] target = shapes[targetIndex];
                    // This prepends 1 to shapes not already equal to nd
                    k = i + target.Length - nd;
                    if (k >= 0)
                    {
                        tmp = target[k];
                        if (tmp == 1)
                        {
                            continue;
                        }

                        if (mit[i] == 1)
                        {
                            mit[i] = tmp;
                        }
                        else if (mit[i] != tmp)
                        {
                            return false;
                        }
                    }
                }
            }

            return true;
        }
    }
}